package kuzha.code.jvm;

import java.io.DataOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.AccessController;
import java.security.PrivilegedAction;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

import kuzha.KuzhaEngine;
import kuzha.ast.Statement;
import kuzha.code.CodeBlock;
import kuzha.code.CodeBuilder;
import kuzha.code.CompileException;

public class JvmCodeBuilder extends CodeBuilder {
	private final KuzhaClassLoader loader;
	private final boolean debugMode;
	private int clsCount = 0;
	
	public JvmCodeBuilder(KuzhaEngine engine, String pname, boolean debugMode) {
		super(engine, pname);
		this.loader = AccessController.doPrivileged(new PrivilegedAction<KuzhaClassLoader>(){
			@Override public KuzhaClassLoader run(){
				return new KuzhaClassLoader();
			}
		});
		this.debugMode = debugMode;
	}
	
	@Override
	public CodeBlock newCodeBlock(Statement node) throws CompileException {
		try {
			return newCodeBlock(node, true);
		}
		catch(CompileFailureException ex) {
			throw new CompileException(ex);
		}
	}
	
	CodeBlock newCodeBlock(Statement node, boolean top) {
		try {
			return (CodeBlock) newCodeBlockClass(node, top).getField("INSTANCE").get(null);
		}
		catch(NoSuchFieldException ex) {
			throw new CompileFailureException(ex);
		}
		catch(IllegalAccessException ex) {
			throw new CompileFailureException(ex);
		}
	}
	Class<?> newCodeBlockClass(Statement node, boolean top) {
		String className = newClassName();
		CodeGenerator generator = new CodeGenerator(this, className, debugMode); 
		byte[] data = generator.generate(node, top);
		loader.putClassData(className, data);
		try {
			return loader.findClass(className);
		}
		catch(ClassNotFoundException ex) {
			throw new CompileFailureException(ex);
		}
	}
	
	private String newClassName() {
		int count;
		synchronized (this) {
			count = clsCount++;
		}
		return getProgramName() + (count == 0 ? "" : "_" + count);
	}
	
	public void saveClassFiles(File dir) throws IOException {
		for(ClassData data: loader.getClassData()) {
			String filepath = data.name + ".class";
			if (dir != null) {
				filepath = dir.getAbsolutePath() + "/" + filepath;
			}
			DataOutputStream out = null;
			try {
				out = new DataOutputStream(new FileOutputStream(filepath));
				out.write(data.data);
			}
			finally {
				if (out != null)
					out.close();
			}
		}
	}
	
	private static class KuzhaClassLoader extends ClassLoader {
		private final Map<String, ClassData> clss = new HashMap<String, ClassData>();
		private final ClassLoader parent;
		
		public KuzhaClassLoader(){
			this.parent = this.getClass().getClassLoader();
		}
		public KuzhaClassLoader(ClassLoader parent) {
			if (parent == null)
				throw new IllegalArgumentException();
			this.parent = parent;
		}
		
		public Collection<ClassData> getClassData() {
			return clss.values();
		}
		public void putClassData(String name, byte[] data) {
			clss.put(name, new ClassData(name, data));
		}
		
		@Override
		protected Class<?> findClass(String name) throws ClassNotFoundException {
			ClassData cd = clss.get(name);
			if (cd != null) {
				synchronized (this) {
					if (cd.cls == null) {
						cd.cls = defineClass(name, cd.data, 0, cd.data.length);
					}
				}
				return cd.cls;
			}
			return parent.loadClass(name);
		}
	}
	private static class ClassData {
		public final String name;
		public final byte[] data;
		public Class<?> cls;
		
		public ClassData(String name, byte[] data) {
			this.name = name;
			this.data = data;
		}
	}
}
